import torch
import random
from copy import deepcopy

class AdvMemory:
    def __init__(self, supplementer, keep_times=20, mem_size=200):
        self.items=[]
        self.supplementer=supplementer
        self.mem_size=mem_size
        self.keep_times=keep_times

    def pack(self, imgs, labels):
        times=[0 for _ in labels]
        imgs=[x for x in imgs]
        return [list(x) for x in zip(deepcopy(imgs), imgs, labels, times)]

    def add(self, imgs, labels):
        self.items.extend(self.pack(imgs, labels))

    def get(self, n):
        idxs=random.sample(range(len(self.items)), n)
        datas=list(zip(*[self.items[x] for x in idxs]))[0:3]
        return idxs, torch.stack(datas[0], dim=0), torch.stack(datas[1], dim=0), list(datas[2])

    def push(self, imgs):
        for i, index in enumerate(sorted(self.idxs, reverse=True)):
            if self.items[index][3]+1>=self.keep_times:
                del self.items[index]
            else:
                self.items[index][1]=imgs[i]
                self.items[index][3]+=1

    def pull(self):
        if self.is_empty():
            self.idxs=list(range(self.supplementer.batch_size*2))
            imgs1, labels1=self.supplementer.next()
            imgs2, labels2=self.supplementer.next()
            imgs=torch.cat((imgs1, imgs2), dim=0)
            ori_imgs=deepcopy(imgs)
            labels=labels1+labels2
            self.add(imgs, labels)
        else:
            if len(self.items) < self.mem_size:
                idxs1 = list(range(len(self.items), len(self.items)+self.supplementer.batch_size))
                imgs1, labels1 = self.supplementer.next()
                ori_imgs1 = deepcopy(imgs1)
                idxs2, ori_imgs2, imgs2, labels2 = self.get(self.supplementer.batch_size)
                self.idxs=idxs1+idxs2
                self.add(imgs1, labels1)
                imgs = torch.cat((imgs1, imgs2), dim=0)
                ori_imgs = torch.cat((ori_imgs1, ori_imgs2), dim=0)
                labels = labels1 + labels2
            else:
                self.idxs, ori_imgs, imgs, labels = self.get(self.supplementer.batch_size*2)

        return ori_imgs, imgs, labels

    def is_empty(self):
        return len(self.items)<=0
